import pandas as pd
import numpy as np
import baostock as bs
import datetime
from tqdm import tqdm
from multiprocessing import Process, Pool 
import os 

class Get_price():
    def __init__(self, path=r'data/stock_data/daily_price', start_date='2001-01-01', freq='d' ):
        self.freq = freq
        self.adjustflag = '2'
        self.start_date = start_date
        self.end_date = str(datetime.date.today())
        self.fields = "date,open,high,low,close,volume,amount,turn"  #,turn,pctChg" #获取数据内容 
        self.path = path 
        if not os.path.exists(self.path):
            os.makedirs(self.path)
        self.index = 'date'
        
    def get_price(self, code_list):
        #### 获取沪深A股历史K线数据 ####
        # 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。“分钟线”不包含指数。
        # 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
        # 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
        for i in tqdm(code_list):
            rs = bs.query_history_k_data_plus(i, fields = self.fields,
                start_date = self.start_date, end_date=self.end_date,
                frequency=self.freq, adjustflag=self.adjustflag)
            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())
            result = pd.DataFrame(data_list, columns=rs.fields)
            result[self.index] = result[self.index]
            result.index = pd.to_datetime(result[self.index])
            result.drop(columns=self.index, inplace=True)
            result.replace('',np.nan, inplace=True)
            result = result.astype(float)  
            result.to_pickle(f'{self.path}/{i}.pkl.gzip')
            # result.to_csv(f'{self.path}/{i}.csv')